package com.black.utils.encrypt.sm2;

import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.params.*;
import org.bouncycastle.math.ec.ECConstants;
import org.bouncycastle.math.ec.ECFieldElement;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.BigIntegers;

import java.math.BigInteger;
import java.security.SecureRandom;

/**
 * Extended SM2 engine for encryption and decryption.
 * 用于加密和解密的扩展SM2引擎。
 */
public class SM2EngineExtend {

    private final Digest digest;
    private boolean forEncryption;
    private ECKeyParameters ecKey;
    private ECDomainParameters ecParams;
    private int curveLength;
    private SecureRandom random;
    private int cipherMode;
    public static int CIPHER_MODE_BC = 0;
    public static int CIPHER_MODE_NORM = 1;

    /**
     * Default constructor using SM3 digest.
     * 使用SM3摘要算法的默认构造函数。
     */
    public SM2EngineExtend() {
        this(new SM3Digest());
    }

    /**
     * Constructor with specified digest.
     * 使用指定摘要算法的构造函数。
     *
     * @param digest The digest algorithm to use.
     *               要使用的摘要算法。
     */
    public SM2EngineExtend(Digest digest) {
        this.digest = digest;
    }

    /**
     * Initialize the engine for encryption or decryption.
     * 初始化引擎以进行加密或解密。
     *
     * @param forEncryption True for encryption, false for decryption.
     *                      加密为true，解密为false。
     * @param param         Cipher parameters.
     *                      加密参数。
     */
    public void init(boolean forEncryption, CipherParameters param) {
        init(forEncryption, CIPHER_MODE_NORM, param);
    }

    /**
     * Initialize the engine for encryption or decryption with specified cipher mode.
     * 使用指定的加密模式初始化引擎以进行加密或解密。
     *
     * @param forEncryption True for encryption, false for decryption.
     *                      加密为true，解密为false。
     * @param cipherMode    The cipher mode to use.
     *                      要使用的加密模式。
     * @param param         Cipher parameters.
     *                      加密参数。
     */
    public void init(boolean forEncryption, int cipherMode, CipherParameters param) {
        this.forEncryption = forEncryption;
        this.cipherMode = cipherMode;
        if (forEncryption) {
            ParametersWithRandom rParam = (ParametersWithRandom) param;

            ecKey = (ECKeyParameters) rParam.getParameters();
            ecParams = ecKey.getParameters();

            ECPoint s = ((ECPublicKeyParameters) ecKey).getQ().multiply(ecParams.getH());
            if (s.isInfinity()) {
                throw new IllegalArgumentException("invalid key: [h]Q at infinity");
            }

            random = rParam.getRandom();
        } else {
            ecKey = (ECKeyParameters) param;
            ecParams = ecKey.getParameters();
        }

        curveLength = (ecParams.getCurve().getFieldSize() + 7) / 8;
    }

    /**
     * Process a block of data for encryption or decryption.
     * 处理一个数据块以进行加密或解密。
     *
     * @param in    The input data.
     *              输入数据。
     * @param inOff The offset in the input data.
     *              输入数据的偏移量。
     * @param inLen The length of the input data.
     *              输入数据的长度。
     * @return The processed data.
     * 处理后的数据。
     * @throws InvalidCipherTextException If an error occurs during processing.
     *                                    如果处理过程中发生错误。
     */
    public byte[] processBlock(byte[] in, int inOff, int inLen) throws InvalidCipherTextException {
        if (forEncryption) {
            return encrypt(in, inOff, inLen);
        } else {
            return decrypt(in, inOff, inLen);
        }
    }

    /**
     * Encrypts the given input data.
     * 加密给定的输入数据。
     *
     * @param in    The input data to be encrypted.
     *              要加密的输入数据。
     * @param inOff The offset in the input data where the data to be encrypted starts.
     *              输入数据中开始加密数据的偏移量。
     * @param inLen The length of the input data to be encrypted.
     *              要加密的输入数据的长度。
     * @return The encrypted data.
     * 加密后的数据。
     * @throws InvalidCipherTextException If an error occurs during encryption.
     *                                    如果加密过程中发生错误。
     */
    private byte[] encrypt(byte[] in, int inOff, int inLen) throws InvalidCipherTextException {
        byte[] c2 = new byte[inLen];

        System.arraycopy(in, inOff, c2, 0, c2.length);

        byte[] c1;
        ECPoint kPB;
        do {
            BigInteger k = nextK();

            ECPoint c1P = ecParams.getG().multiply(k).normalize();

            c1 = c1P.getEncoded(false);

            kPB = ((ECPublicKeyParameters) ecKey).getQ().multiply(k).normalize();

            kdf(digest, kPB, c2);
        }
        while (notEncrypted(c2, in, inOff));

        byte[] c3 = new byte[digest.getDigestSize()];

        addFieldElement(digest, kPB.getAffineXCoord());
        digest.update(in, inOff, inLen);
        addFieldElement(digest, kPB.getAffineYCoord());

        digest.doFinal(c3, 0);
        if (cipherMode == CIPHER_MODE_NORM) {
            return Arrays.concatenate(c1, c3, c2);
        }
        return Arrays.concatenate(c1, c2, c3);
    }

    /**
     * Decrypts the given input data.
     * 解密给定的输入数据。
     *
     * @param in    The input data to be decrypted.
     *              要解密的输入数据。
     * @param inOff The offset in the input data where the data to be decrypted starts.
     *              输入数据中开始解密数据的偏移量。
     * @param inLen The length of the input data to be decrypted.
     *              要解密的输入数据的长度。
     * @return The decrypted data.
     * 解密后的数据。
     * @throws InvalidCipherTextException If an error occurs during decryption.
     *                                    如果解密过程中发生错误。
     */
    private byte[] decrypt(byte[] in, int inOff, int inLen) throws InvalidCipherTextException {
        byte[] c1 = new byte[curveLength * 2 + 1];

        System.arraycopy(in, inOff, c1, 0, c1.length);

        ECPoint c1P = ecParams.getCurve().decodePoint(c1);

        ECPoint s = c1P.multiply(ecParams.getH());
        if (s.isInfinity()) {
            throw new InvalidCipherTextException("[h]C1 at infinity");
        }

        c1P = c1P.multiply(((ECPrivateKeyParameters) ecKey).getD()).normalize();

        byte[] c2 = new byte[inLen - c1.length - digest.getDigestSize()];
        if (cipherMode == CIPHER_MODE_BC) {
            System.arraycopy(in, inOff + c1.length, c2, 0, c2.length);
        } else {
            System.arraycopy(in, inOff + c1.length + digest.getDigestSize(), c2, 0, c2.length);
        }

        kdf(digest, c1P, c2);

        byte[] c3 = new byte[digest.getDigestSize()];

        addFieldElement(digest, c1P.getAffineXCoord());
        digest.update(c2, 0, c2.length);
        addFieldElement(digest, c1P.getAffineYCoord());

        digest.doFinal(c3, 0);

        int check = 0;

        if (cipherMode == CIPHER_MODE_BC) {
            for (int i = 0; i != c3.length; i++) {
                check |= c3[i] ^ in[c1.length + c2.length + i];
            }
        } else {
            for (int i = 0; i != c3.length; i++) {
                check |= c3[i] ^ in[c1.length + i];
            }
        }

        clearBlock(c1);
        clearBlock(c3);

        if (check != 0) {
            clearBlock(c2);
            throw new InvalidCipherTextException("invalid cipher text");
        }

        return c2;
    }

    /**
     * Checks if the encrypted data matches the original input data starting from a specific offset.
     * 检查加密数据是否与从特定偏移量开始的原始输入数据匹配。
     *
     * @param encData The encrypted data to be checked.
     *                要检查的加密数据。
     * @param in      The original input data.
     *                原始输入数据。
     * @param inOff   The offset in the input data where the comparison starts.
     *                输入数据中开始比较的偏移量。
     * @return True if the encrypted data matches the original input data, false otherwise.
     * 如果加密数据与原始输入数据匹配，则返回true，否则返回false。
     */
    private boolean notEncrypted(byte[] encData, byte[] in, int inOff) {
        for (int i = 0; i != encData.length; i++) {
            if (encData[i] != in[inOff]) {
                return false;
            }
        }

        return true;
    }

    /**
     * Key derivation function (KDF) to generate a derived key based on a digest and elliptic curve point.
     * 基于摘要和椭圆曲线点生成派生密钥的密钥派生函数（KDF）。
     *
     * @param digest  The digest algorithm to use.
     *                使用的摘要算法。
     * @param c1      The elliptic curve point.
     *                椭圆曲线点。
     * @param encData The encrypted data to be modified by the derived key.
     *                由派生密钥修改的加密数据。
     */
    private void kdf(Digest digest, ECPoint c1, byte[] encData) {
        int ct = 1;
        int v = digest.getDigestSize();

        byte[] buf = new byte[digest.getDigestSize()];
        int off = 0;

        for (int i = 1; i <= ((encData.length + v - 1) / v); i++) {
            addFieldElement(digest, c1.getAffineXCoord());
            addFieldElement(digest, c1.getAffineYCoord());
            digest.update((byte) (ct >> 24));
            digest.update((byte) (ct >> 16));
            digest.update((byte) (ct >> 8));
            digest.update((byte) ct);

            digest.doFinal(buf, 0);

            if (off + buf.length < encData.length) {
                xor(encData, buf, off, buf.length);
            } else {
                xor(encData, buf, off, encData.length - off);
            }

            off += buf.length;
            ct++;
        }
    }

    /**
     * Performs an XOR operation between the data and the key derivation function (KDF) output.
     * 对数据和密钥派生函数（KDF）输出执行异或操作。
     *
     * @param data       The data to be XORed.
     *                   要进行异或操作的数据。
     * @param kdfOut     The output from the KDF.
     *                   来自KDF的输出。
     * @param dOff       The offset in the data array where the XOR operation starts.
     *                   数据数组中开始执行异或操作的偏移量。
     * @param dRemaining The number of bytes to XOR.
     *                   要进行异或操作的字节数。
     */
    private void xor(byte[] data, byte[] kdfOut, int dOff, int dRemaining) {
        for (int i = 0; i != dRemaining; i++) {
            data[dOff + i] ^= kdfOut[i];
        }
    }

    /**
     * Generates a random integer k suitable for cryptographic operations with the specified bit length.
     * 生成适用于指定位长度的密码操作的随机整数 k。
     *
     * @return A randomly generated BigInteger k.
     * 随机生成的 BigInteger k。
     */
    private BigInteger nextK() {
        int qBitLength = ecParams.getN().bitLength();

        BigInteger k;
        do {
            k = new BigInteger(qBitLength, random);
        }
        while (k.equals(ECConstants.ZERO) || k.compareTo(ecParams.getN()) >= 0);

        return k;
    }

    /**
     * Adds the specified elliptic curve field element v to the digest for cryptographic processing.
     * 将指定的椭圆曲线域元素 v 添加到用于密码处理的摘要中。
     *
     * @param digest The digest algorithm instance to update.
     *               要更新的摘要算法实例。
     * @param v      The elliptic curve field element to add to the digest.
     *               要添加到摘要的椭圆曲线域元素。
     */
    private void addFieldElement(Digest digest, ECFieldElement v) {
        byte[] p = BigIntegers.asUnsignedByteArray(curveLength, v.toBigInteger());

        digest.update(p, 0, p.length);
    }

    /**
     * Clears the contents of the specified byte array block by setting all elements to zero.
     * 通过将所有元素设置为零来清除指定的字节数组块的内容。
     *
     * @param block The byte array block to be cleared.
     *              要清除的字节数组块。
     */
    private void clearBlock(byte[] block) {
        Arrays.fill(block, (byte) 0);
    }

}
